import torch
import math
import torch.nn.functional as F
import numpy as np
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from .utils import conv, update_registered_buffers, deconv

import torch.nn as nn
from compressai.entropy_models import EntropyBottleneck, GaussianConditional
from compressai.ans import BufferedRansEncoder, RansDecoder
from compressai.layers import conv3x3, subpel_conv3x3
from compressai.ops import ste_round
from .base import CompressionModel

from functools import partial

from timm.models.vision_transformer import PatchEmbed, Block

from .util.pos_embed import get_2d_sincos_pos_embed


class MaskedAutoencoderViT(CompressionModel):
    """ Masked Autoencoder with VisionTransformer backbone
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 embed_dim=1024, depth=24, num_heads=16,
                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False, num_slices=32, frozen_stages=-1, encoder_embed_dim_out=138, entory_embed_dim_in=128, decoder_embed_dim_in=138, decoder_embed_dim_out=196):
        super().__init__()

        self.in_chans = in_chans
        self.encoder_embed_dim_out = encoder_embed_dim_out
        self.entory_embed_dim_in = entory_embed_dim_in
        self.decoder_embed_dim_in = decoder_embed_dim_in
        self.decoder_embed_dim_out = decoder_embed_dim_out
        self.embed_dim = embed_dim
        self.frozen_stages = frozen_stages
        self.num_slices = num_slices
        self.max_support_slices = num_slices // 2
        self.img_size = img_size

        # --------------------------------------------------------------------------
        # MAE encoder specifics
        self.patch_embed = PatchEmbed(
            img_size, patch_size, in_chans, embed_dim)

        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding

        self.blocks = nn.ModuleList([
            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True,
                norm_layer=norm_layer)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)
        self.conv2d_down = nn.Conv2d(
            in_channels=self.encoder_embed_dim_out, out_channels=self.entory_embed_dim_in, kernel_size=1)
        self.conv2d_up = nn.Conv2d(
            in_channels=self.entory_embed_dim_in, out_channels=self.decoder_embed_dim_in, kernel_size=1)
        # --------------------------------------------------------------------------

        # --------------------------------------------------------------------------
        # MAE decoder specifics
        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)

        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))

        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding

        self.decoder_blocks = nn.ModuleList([
            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio,
                  qkv_bias=True, norm_layer=norm_layer)
            for i in range(decoder_depth)])

        self.decoder_norm = norm_layer(decoder_embed_dim)
        self.decoder_pred = nn.Linear(
            decoder_embed_dim, patch_size**2 * in_chans, bias=True)  # decoder to patch
        self.end_conv = nn.Sequential(nn.Conv2d(num_patches, num_patches * patch_size ** 2, kernel_size=5, stride=1, padding=2),
                                      nn.PixelShuffle(patch_size),
                                      nn.Conv2d(
                                          num_patches, 3, kernel_size=3, stride=1, padding=1),
                                      )
        # --------------------------------------------------------------------------

        self.norm_pix_loss = norm_pix_loss

        self.initialize_weights()

        self.g_a = None
        self.g_s = None

        self.h_a = nn.Sequential(
            conv3x3(128, 128),
            nn.GELU(),
            conv3x3(128, 112),
            nn.GELU(),
            conv3x3(112, 96, stride=2),
            nn.GELU(),
            conv3x3(96, 80),
            nn.GELU(),
            conv3x3(80, 64, stride=2),
        )

        self.h_mean_s = nn.Sequential(
            conv3x3(64, 80),
            nn.GELU(),
            subpel_conv3x3(80, 96, 2),
            nn.GELU(),
            conv3x3(96, 112),
            nn.GELU(),
            subpel_conv3x3(112, 128, 2),
            nn.GELU(),
            conv3x3(128, 128),
        )
        self.h_scale_s = nn.Sequential(
            conv3x3(64, 80),
            nn.GELU(),
            subpel_conv3x3(80, 96, 2),
            nn.GELU(),
            conv3x3(96, 112),
            nn.GELU(),
            subpel_conv3x3(112, 128, 2),
            nn.GELU(),
            conv3x3(128, 128),
        )

        # 不确定这么改对不对

        self.cc_mean_transforms = nn.ModuleList(
            nn.Sequential(
                conv(128 + 4 * min(i, 16), 64, stride=1, kernel_size=3),
                nn.GELU(),
                conv(64, 32, stride=1, kernel_size=3),
                nn.GELU(),
                conv(32, 16, stride=1, kernel_size=3),
                nn.GELU(),
                conv(16, 8, stride=1, kernel_size=3),
                nn.GELU(),
                conv(8, 4, stride=1, kernel_size=3),
            ) for i in range(num_slices)
        )
        self.cc_scale_transforms = nn.ModuleList(
            nn.Sequential(
                conv(128 + 4 * min(i, 16), 64, stride=1, kernel_size=3),
                nn.GELU(),
                conv(64, 32, stride=1, kernel_size=3),
                nn.GELU(),
                conv(32, 16, stride=1, kernel_size=3),
                nn.GELU(),
                conv(16, 8, stride=1, kernel_size=3),
                nn.GELU(),
                conv(8, 4, stride=1, kernel_size=3),
            ) for i in range(num_slices)
        )
        self.lrp_transforms = nn.ModuleList(
            nn.Sequential(
                conv(128 + 4 * min(i + 1, 17), 64, stride=1, kernel_size=3),
                nn.GELU(),
                conv(64, 32, stride=1, kernel_size=3),
                nn.GELU(),
                conv(32, 16, stride=1, kernel_size=3),
                nn.GELU(),
                conv(16, 8, stride=1, kernel_size=3),
                nn.GELU(),
                conv(8, 4, stride=1, kernel_size=3),
            ) for i in range(num_slices)
        )
        self.entropy_bottleneck = EntropyBottleneck(
            self.entory_embed_dim_in//2)
        self.gaussian_conditional = GaussianConditional(None)
        self._freeze_stages()

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.patch_embed.eval()
            for param in self.patch_embed.parameters():
                param.requires_grad = False

        if self.frozen_stages >= 1 and self.ape:
            self.absolute_pos_embed.requires_grad = False

        if self.frozen_stages >= 2:
            self.pos_drop.eval()
            for i in range(0, self.frozen_stages - 1):
                m = self.layers[i]
                m.eval()
                for param in m.parameters():
                    param.requires_grad = False

    def initialize_weights(self):
        # initialization
        # initialize (and freeze) pos_embed by sin-cos embedding
        pos_embed = get_2d_sincos_pos_embed(
            self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)

        self.pos_embed.data.copy_(
            torch.from_numpy(pos_embed).float().unsqueeze(0))

        decoder_pos_embed = get_2d_sincos_pos_embed(
            self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
        self.decoder_pos_embed.data.copy_(
            torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))

        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
        torch.nn.init.normal_(self.cls_token, std=.02)
        torch.nn.init.normal_(self.mask_token, std=.02)

        # initialize nn.Linear and nn.LayerNorm
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def patchify(self, imgs):
        """
        imgs: (N, 3, H, W)
        x: (N, L, patch_size**2 *3)
        """
        p = self.patch_embed.patch_size[0]
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0

        h = w = imgs.shape[2] // p
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
        x = torch.einsum('nchpwq->nhwpqc', x)
        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
        return x

    def unpatchify(self, x):
        """
        x: (N, L, patch_size**2 *3)
        imgs: (N, 3, H, W)
        """
        p = self.patch_embed.patch_size[0]
        h = w = int(x.shape[1]**.5)
        assert h * w == x.shape[1]

        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        return imgs

    def random_masking(self, x, mask_ratio):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        N, L, D = x.shape  # batch, length, dim
        len_keep = int(L * (1 - mask_ratio))

        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]

        # sort noise for each sample
        # ascend: small is keep, large is remove
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        # keep the first subset
        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(
            x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

    # def forward_encoder(self, x, mask_ratio):
    #     # embed patches
    #     x = self.patch_embed(x)

    #     # add pos embed w/o cls token
    #     x = x + self.pos_embed[:, 1:, :]

    #     # masking: length -> length * mask_ratio
    #     x, mask, ids_restore = self.random_masking(x, mask_ratio)

    #     # append cls token
    #     # cls_token = self.cls_token + self.pos_embed[:, :1, :]
    #     # cls_tokens = cls_token.expand(x.shape[0], -1, -1)
    #     # x = torch.cat((cls_tokens, x), dim=1)

    #     # apply Transformer blocks
    #     for blk in self.blocks:
    #         x = blk(x)
    #     x = self.norm(x)

    #     return x, mask, ids_restore

    def forward_decoder(self, x, ids_restore):
        # embed tokens
        x = self.decoder_embed(x)

        # append mask tokens to sequence
        mask_tokens = self.mask_token.repeat(
            x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
        x_ = torch.cat([x[:, :, :], mask_tokens], dim=1)  # no cls token
        x_ = torch.gather(
            x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token

        # add pos embed
        x = x + self.decoder_pos_embed

        # apply Transformer blocks
        for blk in self.decoder_blocks:
            x = blk(x)
        x = self.decoder_norm(x)

        # predictor projection
        x = self.decoder_pred(x)

        # remove cls token
        x = x[:, 1:, :]

        return x

    def forward_loss(self, imgs, pred, mask):
        """
        imgs: [N, 3, H, W]
        pred: [N, L, p*p*3]
        mask: [N, L], 0 is keep, 1 is remove, 
        """
        target = self.patchify(imgs)
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.e-6)**.5

        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch

        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
        return loss

    def forward(self, imgs, mask_ratio=0.3):
       # ---------------------mae的forward_encoder--------------------------
        # embed patches
        x = self.patch_embed(imgs)

        ''' 
        print(x.shape)  [16, 196, 1024]  [bs,L,P**2*3] ? 
        '''

        # Wh, Ww = x.size(2), x.size(3)

        # add pos embed w/o cls token
        x = x + self.pos_embed[:, 1:, :]

        # masking: length -> length * mask_ratio
        x, mask, ids_restore = self.random_masking(x, mask_ratio)

        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # apply Transformer blocks
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        # x.shape == torch.Size([16, 49, 1024])
        # ---------------------mae的forward_encoder--------------------------

        y = x
        C = self.entory_embed_dim_in
        Wh, Ww = 32, 32
        # self.entropy_embed = torch.nn.Linear(
        #     self.encoder_embed_dim_out, self.entory_embed_dim_in, bias=True)
        # y = self.entropy_embed(y)

        # y = torch.nn.Linear(y.shape[0], self.encoder_embed_dim_out*y.shape[2],
        #                     self.entory_embed_dim_in, bias=True)(x)

        # y = y.view(-1, self.entory_embed_dim_in, Wh, Ww)

        y = y.view(y.shape[0], y.shape[1], Wh, Ww)
        features_encoder_out = y
        y = self.conv2d_down(y)
        y_shape = y.shape[2:]

        z = self.h_a(y)
        _, z_likelihoods = self.entropy_bottleneck(z)
        z_offset = self.entropy_bottleneck._get_medians()
        z_tmp = z - z_offset
        z_hat = ste_round(z_tmp) + z_offset

        latent_scales = self.h_scale_s(z_hat)
        latent_means = self.h_mean_s(z_hat)

        y_slices = y.chunk(self.num_slices, 1)
        y_hat_slices = []
        y_likelihood = []

        for slice_index, y_slice in enumerate(y_slices):
            support_slices = (y_hat_slices if self.max_support_slices <
                              0 else y_hat_slices[:self.max_support_slices])
            mean_support = torch.cat([latent_means] + support_slices, dim=1)
            mu = self.cc_mean_transforms[slice_index](mean_support)
            mu = mu[:, :, :y_shape[0], :y_shape[1]]

            scale_support = torch.cat([latent_scales] + support_slices, dim=1)
            scale = self.cc_scale_transforms[slice_index](scale_support)
            scale = scale[:, :, :y_shape[0], :y_shape[1]]

            _, y_slice_likelihood = self.gaussian_conditional(
                y_slice, scale, mu)

            y_likelihood.append(y_slice_likelihood)
            y_hat_slice = ste_round(y_slice - mu) + mu

            lrp_support = torch.cat([mean_support, y_hat_slice], dim=1)
            lrp = self.lrp_transforms[slice_index](lrp_support)
            lrp = 0.5 * torch.tanh(lrp)
            y_hat_slice += lrp

            y_hat_slices.append(y_hat_slice)

        y_hat = torch.cat(y_hat_slices, dim=1)
        y_likelihoods = torch.cat(y_likelihood, dim=1)
        y_hat = self.conv2d_up(y_hat)
        features_entroy_out = y_hat
        y_hat = y_hat.view(y_hat.shape[0], -1, Wh*Ww)
        # y_hat = y_hat.permute(0, 2, 3, 1).contiguous().view(-1, Wh*Ww, self.decoder_embed_dim_in)

        # for i in range(self.num_layers):
        #     layer = self.syn_layers[i]
        #     y_hat, Wh, Ww = layer(y_hat, Wh, Ww)

        # ---------------------mae的forward_decoder--------------------------
        # num_masked = int(mask_ratio * num_patches)
        # mask_tokens = self.mask_token[, , :].repeat(y_hat.shape[0], num_masked, 1)
        y_rec = self.forward_decoder(y_hat, ids_restore)
        x_hat = self.unpatchify(y_rec)
        # y_rec.shape=[16,196,768]
        # x_hat = y_rec.view(-1, self.in_chans, self.img_size, self.img_size)

        # x_hat = self.end_conv(
        #     y_rec.view(-1, Wh, self.img_size, self.decoder_embed_dim_out).permute(0, 3, 1, 2).contiguous())

        return {
            "x_hat": x_hat,
            "likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
            "features_encoder_out": features_encoder_out,
            "features_entroy_out": features_encoder_out,
        }
        # pred = self.forward_decoder(latent, ids_restore)  # [N, L, p*p*3]
        # loss = self.forward_loss(imgs, pred, mask)
        # return loss, pred, mask

    def update(self, scale_table=None, force=False):
        if scale_table is None:
            scale_table = get_scale_table()
        updated = self.gaussian_conditional.update_scale_table(
            scale_table, force=force)
        updated |= super().update(force=force)
        return updated

    def load_state_dict(self, state_dict, strict_value):
        # update_registered_buffers(
        #     self.gaussian_conditional,
        #     "gaussian_conditional",
        #     ["_quantized_cdf", "_offset", "_cdf_length", "scale_table"],
        #     state_dict,
        # )
        super().load_state_dict(state_dict, strict_value)

    @classmethod
    def from_state_dict(cls, state_dict):
        """Return a new model instance from `state_dict`."""
        net = cls()
        net.load_state_dict(state_dict)
        return net

    def compress(self, x):
        # encoder部分--------------------------------------------------------------

        # [1, 3, 512, 512]------------------------
        # print(x.shape)
        # x = self.patch_embed(x)
        # # [1, 48, 256, 256]----------------------
        # print(x.shape)

        # Wh, Ww = x.size(2), x.size(3)

        # ''' flatten把第二维及之后的维度都压缩到第二维
        # x.flatten(2).shape == [1,48,65536]
        # transpose交换第一维和第二维
        # x.flatten(2).transpose(1, 2).shape == [1,65536,48]
        # 最终达到的效果是h和w维度压缩成一维,通道维度放在最后p0
        # '''
        # x = x.flatten(2).transpose(1, 2)
        # # [1, 65536, 48]-------------------------
        # print(x.shape)
        # for i in range(self.num_layers):
        #     layer = self.layers[i]
        #     x, Wh, Ww = layer(x, Wh, Ww)
        mask_ratio = 0.3
        x = forward_encoder(self, x, mask_ratio)
        Wh, Ww = x.size(2), x.size(3)
        y = x
        # [1, 1024, 384]---------------------------
        print(y.shape)

        # encoder部分--------------------------------------------------------------------

        C = self.embed_dim * 8
        y = y.view(-1, Wh, Ww, C).permute(0, 3, 1, 2).contiguous()
        # [1, 384, 32, 32]-------------------------
        print(y.shape)
        y_shape = y.shape[2:]

        z = self.h_a(y)
        z_strings = self.entropy_bottleneck.compress(z)
        z_hat = self.entropy_bottleneck.decompress(z_strings, z.size()[-2:])

        latent_scales = self.h_scale_s(z_hat)
        latent_means = self.h_mean_s(z_hat)

        y_slices = y.chunk(self.num_slices, 1)
        y_hat_slices = []

        cdf = self.gaussian_conditional.quantized_cdf.tolist()
        cdf_lengths = self.gaussian_conditional.cdf_length.reshape(
            -1).int().tolist()
        offsets = self.gaussian_conditional.offset.reshape(-1).int().tolist()

        encoder = BufferedRansEncoder()
        symbols_list = []
        indexes_list = []
        y_strings = []

        for slice_index, y_slice in enumerate(y_slices):
            support_slices = (y_hat_slices if self.max_support_slices <
                              0 else y_hat_slices[:self.max_support_slices])

            mean_support = torch.cat([latent_means] + support_slices, dim=1)
            mu = self.cc_mean_transforms[slice_index](mean_support)
            mu = mu[:, :, :y_shape[0], :y_shape[1]]

            scale_support = torch.cat([latent_scales] + support_slices, dim=1)
            scale = self.cc_scale_transforms[slice_index](scale_support)
            scale = scale[:, :, :y_shape[0], :y_shape[1]]

            index = self.gaussian_conditional.build_indexes(scale)
            y_q_slice = self.gaussian_conditional.quantize(
                y_slice, "symbols", mu)
            y_hat_slice = y_q_slice + mu

            symbols_list.extend(y_q_slice.reshape(-1).tolist())
            indexes_list.extend(index.reshape(-1).tolist())

            lrp_support = torch.cat([mean_support, y_hat_slice], dim=1)
            lrp = self.lrp_transforms[slice_index](lrp_support)
            lrp = 0.5 * torch.tanh(lrp)
            y_hat_slice += lrp

            y_hat_slices.append(y_hat_slice)
        encoder.encode_with_indexes(
            symbols_list, indexes_list, cdf, cdf_lengths, offsets)

        y_string = encoder.flush()
        y_strings.append(y_string)
        return {"strings": [y_strings, z_strings], "shape": z.size()[-2:]}

    def decompress(self, strings, shape):
        assert isinstance(strings, list) and len(strings) == 2
        z_hat = self.entropy_bottleneck.decompress(strings[1], shape)
        latent_scales = self.h_scale_s(z_hat)
        latent_means = self.h_mean_s(z_hat)

        y_shape = [z_hat.shape[2] * 4, z_hat.shape[3] * 4]
        Wh, Ww = y_shape
        C = self.embed_dim * 8

        y_string = strings[0][0]

        y_hat_slices = []
        cdf = self.gaussian_conditional.quantized_cdf.tolist()
        cdf_lengths = self.gaussian_conditional.cdf_length.reshape(
            -1).int().tolist()
        offsets = self.gaussian_conditional.offset.reshape(-1).int().tolist()

        decoder = RansDecoder()
        decoder.set_stream(y_string)

        for slice_index in range(self.num_slices):
            support_slices = (y_hat_slices if self.max_support_slices <
                              0 else y_hat_slices[:self.max_support_slices])
            mean_support = torch.cat([latent_means] + support_slices, dim=1)
            mu = self.cc_mean_transforms[slice_index](mean_support)
            mu = mu[:, :, :y_shape[0], :y_shape[1]]

            scale_support = torch.cat([latent_scales] + support_slices, dim=1)
            scale = self.cc_scale_transforms[slice_index](scale_support)
            scale = scale[:, :, :y_shape[0], :y_shape[1]]

            index = self.gaussian_conditional.build_indexes(scale)

            rv = decoder.decode_stream(
                index.reshape(-1).tolist(), cdf, cdf_lengths, offsets)
            rv = torch.Tensor(rv).reshape(1, -1, y_shape[0], y_shape[1])
            y_hat_slice = self.gaussian_conditional.dequantize(rv, mu)

            lrp_support = torch.cat([mean_support, y_hat_slice], dim=1)
            lrp = self.lrp_transforms[slice_index](lrp_support)
            lrp = 0.5 * torch.tanh(lrp)
            y_hat_slice += lrp

            y_hat_slices.append(y_hat_slice)

        y_hat = torch.cat(y_hat_slices, dim=1)
        y_hat = y_hat.permute(0, 2, 3, 1).contiguous().view(-1, Wh * Ww, C)
        for i in range(self.num_layers):
            layer = self.syn_layers[i]
            y_hat, Wh, Ww = layer(y_hat, Wh, Ww)

        x_hat = self.end_conv(
            y_hat.view(-1, Wh, Ww, self.decoder_embed_dim).permute(0, 3, 1, 2).contiguous()).clamp_(0, 1)
        return {"x_hat": x_hat}


def mae_vit_base_patch16_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=16, embed_dim=768, depth=12, num_heads=12,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def mae_vit_large_patch16_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def mae_vit_huge_patch14_dec512d8b(**kwargs):
    model = MaskedAutoencoderViT(
        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


# set recommended archs
mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b  # decoder: 512 dim, 8 blocks
